import pickle
import os
import glob
import collections
import numpy as np


train2014 = "labels/train2014"
val2014 = "labels/val2014"


labels = collections.defaultdict(list)

for f in glob.glob(train2014+"/*.txt"):
    name = os.path.basename(os.path.abspath(f))[:-4]
    lines = open(f, "r").readlines()
    coor = []
    for line in lines:
        coor.append([float(c) for c in line.strip().split(" ")])
    coor = np.array(coor).reshape(-1, 5)
    labels["train"].append({name:coor})

for f in glob.glob(val2014+"/*.txt"):
    name = os.path.basename(os.path.abspath(f))[:-4]
    lines = open(f, "r").readlines()
    coor = []
    for line in lines:
        coor.append([float(c) for c in line.strip().split(" ")])
    coor = np.array(coor).reshape(-1, 5)
    labels["val"].append({name:coor})

with open("coco2014_lables.pkl", "wb") as f:
    pickle.dump(labels,f)
